{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "matplotlib.get_backend :  module://ipykernel.pylab.backend_inline\n"
     ]
    }
   ],
   "source": [
    "from create_parquet import *\n",
    "from data import *\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\") "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build node frame from graphs "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_node_from_graph(molecule_file): \n",
    "    '''\n",
    "      - molecule file:  path to %molecule_name.pickle\n",
    "    Returns: \n",
    "      Convert the pickled graph to a padded vector with all the molecule information \n",
    "    '''\n",
    "    molecule_name = molecule_file.split('/')[-1].strip('.pickle')\n",
    "    graph = read_pickle_from_file(molecule_file)\n",
    "    molecule_name = graph.molecule_name\n",
    "    node_feats = np.concatenate(graph.node,-1)\n",
    "    num_node, node_dim = node_feats.shape \n",
    "    node = pd.DataFrame(node_feats)\n",
    "    node.columns = ['symbol', 'acceptor', 'donor', 'aromatic', 'hybridization', 'num_h', 'atomic',]\n",
    "    node['num_nodes'] = num_node\n",
    "    node['node_dim'] = node_dim\n",
    "    node['molecule_name'] = molecule_name\n",
    "    node['atom_index'] = list(range(num_node))\n",
    "    return  node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>symbol</th>\n",
       "      <th>acceptor</th>\n",
       "      <th>donor</th>\n",
       "      <th>aromatic</th>\n",
       "      <th>hybridization</th>\n",
       "      <th>num_h</th>\n",
       "      <th>atomic</th>\n",
       "      <th>num_nodes</th>\n",
       "      <th>node_dim</th>\n",
       "      <th>molecule_name</th>\n",
       "      <th>atom_index</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>7.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>4.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_133885</td>\n",
       "      <td>15</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    symbol  acceptor  donor  aromatic  hybridization  num_h  atomic  \\\n",
       "0      2.0       0.0    0.0       0.0            3.0    2.0     6.0   \n",
       "1      3.0       0.0    1.0       0.0            3.0    0.0     7.0   \n",
       "2      2.0       0.0    0.0       0.0            3.0    1.0     6.0   \n",
       "3      2.0       0.0    0.0       0.0            3.0    1.0     6.0   \n",
       "4      2.0       0.0    0.0       0.0            3.0    1.0     6.0   \n",
       "5      4.0       1.0    0.0       0.0            3.0    0.0     8.0   \n",
       "6      2.0       0.0    0.0       0.0            3.0    0.0     6.0   \n",
       "7      2.0       0.0    0.0       0.0            3.0    1.0     6.0   \n",
       "8      2.0       0.0    0.0       0.0            3.0    1.0     6.0   \n",
       "9      1.0       0.0    0.0       0.0            0.0    0.0     1.0   \n",
       "10     1.0       0.0    0.0       0.0            0.0    0.0     1.0   \n",
       "11     1.0       0.0    0.0       0.0            0.0    0.0     1.0   \n",
       "12     1.0       0.0    0.0       0.0            0.0    0.0     1.0   \n",
       "13     1.0       0.0    0.0       0.0            0.0    0.0     1.0   \n",
       "14     1.0       0.0    0.0       0.0            0.0    0.0     1.0   \n",
       "15     1.0       0.0    0.0       0.0            0.0    0.0     1.0   \n",
       "\n",
       "    num_nodes  node_dim     molecule_name  atom_index  \n",
       "0          16         7  dsgdb9nsd_133885           0  \n",
       "1          16         7  dsgdb9nsd_133885           1  \n",
       "2          16         7  dsgdb9nsd_133885           2  \n",
       "3          16         7  dsgdb9nsd_133885           3  \n",
       "4          16         7  dsgdb9nsd_133885           4  \n",
       "5          16         7  dsgdb9nsd_133885           5  \n",
       "6          16         7  dsgdb9nsd_133885           6  \n",
       "7          16         7  dsgdb9nsd_133885           7  \n",
       "8          16         7  dsgdb9nsd_133885           8  \n",
       "9          16         7  dsgdb9nsd_133885           9  \n",
       "10         16         7  dsgdb9nsd_133885          10  \n",
       "11         16         7  dsgdb9nsd_133885          11  \n",
       "12         16         7  dsgdb9nsd_133885          12  \n",
       "13         16         7  dsgdb9nsd_133885          13  \n",
       "14         16         7  dsgdb9nsd_133885          14  \n",
       "15         16         7  dsgdb9nsd_133885          15  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_node_from_graph('/rapids/notebooks/srabhi/champs-2019/input/structure/graph2/dsgdb9nsd_133885.pickle')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 131k/131k [02:24<00:00, 905it/s]    \n",
      "130772it [00:00, 411905.84it/s]\n"
     ]
    }
   ],
   "source": [
    "from parallel_process import parallel_process\n",
    "files = glob.glob('/rapids/notebooks/srabhi/champs-2019/input/structure/graph2/*.pickle')\n",
    "frames = parallel_process(files, get_node_from_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_frame = pd.concat(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_frame.to_csv('/rapids/notebooks/srabhi/champs-2019/input/parquet/baseline_node_frame.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>symbol</th>\n",
       "      <th>acceptor</th>\n",
       "      <th>donor</th>\n",
       "      <th>aromatic</th>\n",
       "      <th>hybridization</th>\n",
       "      <th>num_h</th>\n",
       "      <th>atomic</th>\n",
       "      <th>num_nodes</th>\n",
       "      <th>node_dim</th>\n",
       "      <th>molecule_name</th>\n",
       "      <th>atom_index</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>21</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_101594</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>7.0</td>\n",
       "      <td>21</td>\n",
       "      <td>7</td>\n",
       "      <td>dsgdb9nsd_101594</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   symbol  acceptor  donor  aromatic  hybridization  num_h  atomic  num_nodes  \\\n",
       "0     2.0       0.0    0.0       0.0            3.0    3.0     6.0         21   \n",
       "1     3.0       0.0    1.0       0.0            2.0    1.0     7.0         21   \n",
       "\n",
       "   node_dim     molecule_name  atom_index  \n",
       "0         7  dsgdb9nsd_101594           0  \n",
       "1         7  dsgdb9nsd_101594           1  "
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_frame.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build coupling dataframe from graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_coupling_from_graph(molecule_file): \n",
    "    '''\n",
    "      - molecule file:  path to %molecule_name.pickle\n",
    "    Returns: \n",
    "      Convert the pickled graph to a padded vector with all the molecule information \n",
    "    '''\n",
    "    molecule_name = molecule_file.split('/')[-1].strip('.pickle')\n",
    "    graph = read_pickle_from_file(molecule_file)\n",
    "    molecule_name = graph.molecule_name\n",
    "    \n",
    "    coupling_feats = np.concatenate([graph.coupling.index, graph.coupling.type.reshape(-1, 1), \n",
    "                               graph.coupling.value.reshape(-1,1), graph.coupling.contribution,\n",
    "                               graph.coupling.id.reshape(-1,1)], -1)\n",
    "\n",
    "    num_coupling, coupling_dim = coupling_feats.shape\n",
    "\n",
    "    #change to cudf\n",
    "    coupling = pd.DataFrame(coupling_feats)\n",
    "    coupling.columns = ['atom_index_0', 'atom_index_1', 'coupling_type', 'scalar_coupling', 'fc', 'sd', 'pso', 'dso', 'id']\n",
    "    coupling['num_coupling'] = num_coupling\n",
    "    coupling['coupling_dim'] = coupling_dim\n",
    "    coupling['molecule_name'] = molecule_name\n",
    "    return  coupling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>atom_index_0</th>\n",
       "      <th>atom_index_1</th>\n",
       "      <th>coupling_type</th>\n",
       "      <th>scalar_coupling</th>\n",
       "      <th>fc</th>\n",
       "      <th>sd</th>\n",
       "      <th>pso</th>\n",
       "      <th>dso</th>\n",
       "      <th>id</th>\n",
       "      <th>num_coupling</th>\n",
       "      <th>coupling_dim</th>\n",
       "      <th>molecule_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>9.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>82.31910</td>\n",
       "      <td>80.45880</td>\n",
       "      <td>0.170039</td>\n",
       "      <td>0.892564</td>\n",
       "      <td>0.797695</td>\n",
       "      <td>3622270.0</td>\n",
       "      <td>72</td>\n",
       "      <td>9</td>\n",
       "      <td>dsgdb9nsd_103915</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>9.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>-2.13186</td>\n",
       "      <td>-2.26579</td>\n",
       "      <td>0.091737</td>\n",
       "      <td>-0.046631</td>\n",
       "      <td>0.088824</td>\n",
       "      <td>3622271.0</td>\n",
       "      <td>72</td>\n",
       "      <td>9</td>\n",
       "      <td>dsgdb9nsd_103915</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   atom_index_0  atom_index_1  coupling_type  scalar_coupling        fc  \\\n",
       "0           9.0           0.0            0.0         82.31910  80.45880   \n",
       "1           9.0           1.0            1.0         -2.13186  -2.26579   \n",
       "\n",
       "         sd       pso       dso         id  num_coupling  coupling_dim  \\\n",
       "0  0.170039  0.892564  0.797695  3622270.0            72             9   \n",
       "1  0.091737 -0.046631  0.088824  3622271.0            72             9   \n",
       "\n",
       "      molecule_name  \n",
       "0  dsgdb9nsd_103915  \n",
       "1  dsgdb9nsd_103915  "
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_coupling_from_graph('/rapids/notebooks/srabhi/champs-2019/input/structure/graph2/dsgdb9nsd_103915.pickle').head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 131k/131k [02:17<00:00, 950it/s]    \n",
      "130772it [00:00, 332063.91it/s]\n"
     ]
    }
   ],
   "source": [
    "from parallel_process import parallel_process\n",
    "files = glob.glob('/rapids/notebooks/srabhi/champs-2019/input/structure/graph2/*.pickle')\n",
    "frames = parallel_process(files, get_coupling_from_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "coupling_frame = pd.concat(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>atom_index_0</th>\n",
       "      <th>atom_index_1</th>\n",
       "      <th>coupling_type</th>\n",
       "      <th>scalar_coupling</th>\n",
       "      <th>fc</th>\n",
       "      <th>sd</th>\n",
       "      <th>pso</th>\n",
       "      <th>dso</th>\n",
       "      <th>id</th>\n",
       "      <th>num_coupling</th>\n",
       "      <th>coupling_dim</th>\n",
       "      <th>molecule_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>9.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>84.505600</td>\n",
       "      <td>83.117900</td>\n",
       "      <td>0.133123</td>\n",
       "      <td>0.548422</td>\n",
       "      <td>0.706185</td>\n",
       "      <td>3528934.0</td>\n",
       "      <td>64</td>\n",
       "      <td>9</td>\n",
       "      <td>dsgdb9nsd_101594</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>9.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>-0.420380</td>\n",
       "      <td>-0.454646</td>\n",
       "      <td>-0.012691</td>\n",
       "      <td>0.098126</td>\n",
       "      <td>-0.051169</td>\n",
       "      <td>3528935.0</td>\n",
       "      <td>64</td>\n",
       "      <td>9</td>\n",
       "      <td>dsgdb9nsd_101594</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>9.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>-0.017699</td>\n",
       "      <td>0.076249</td>\n",
       "      <td>0.003486</td>\n",
       "      <td>0.026211</td>\n",
       "      <td>-0.123645</td>\n",
       "      <td>3528936.0</td>\n",
       "      <td>64</td>\n",
       "      <td>9</td>\n",
       "      <td>dsgdb9nsd_101594</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   atom_index_0  atom_index_1  coupling_type  scalar_coupling         fc  \\\n",
       "0           9.0           0.0            0.0        84.505600  83.117900   \n",
       "1           9.0           1.0            4.0        -0.420380  -0.454646   \n",
       "2           9.0           2.0            2.0        -0.017699   0.076249   \n",
       "\n",
       "         sd       pso       dso         id  num_coupling  coupling_dim  \\\n",
       "0  0.133123  0.548422  0.706185  3528934.0            64             9   \n",
       "1 -0.012691  0.098126 -0.051169  3528935.0            64             9   \n",
       "2  0.003486  0.026211 -0.123645  3528936.0            64             9   \n",
       "\n",
       "      molecule_name  \n",
       "0  dsgdb9nsd_101594  \n",
       "1  dsgdb9nsd_101594  \n",
       "2  dsgdb9nsd_101594  "
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coupling_frame.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "coupling_frame.to_csv('/rapids/notebooks/srabhi/champs-2019/input/parquet/baseline_coupling_frame.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build edge frame from graph "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data import *\n",
    "def get_edge_from_graph(molecule_file): \n",
    "    '''\n",
    "      - molecule file:  path to %molecule_name.pickle\n",
    "    Returns: \n",
    "      Convert the pickled graph to a padded vector with all the molecule information \n",
    "    '''\n",
    "    molecule_name = molecule_file.split('/')[-1].strip('.pickle')\n",
    "    graph = read_pickle_from_file(molecule_file)\n",
    "    molecule_name = graph.molecule_name\n",
    "    edge_feats = np.concatenate(graph.edge,-1)\n",
    "    edge_feats = np.concatenate([graph.edge_index, edge_feats], -1)\n",
    "    num_edge, edge_dim = edge_feats.shape \n",
    "    infor = [molecule_name, num_edge, edge_dim]\n",
    "    edge = pd.DataFrame(edge_feats)\n",
    "    edge.columns = ['atom_index_0', 'atom_index_1', 'edge_type', 'distance', 'angle']\n",
    "    edge['molecule_name'] = molecule_name\n",
    "    edge['num_edge'] = num_edge\n",
    "    edge['edge_dim'] = edge_dim\n",
    "    return  edge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = glob.glob('/rapids/notebooks/srabhi/champs-2019/input/structure/graph2/*.pickle')\n",
    "#t = get_edge_from_graph(molecule_file+molecule_name+'.pickle')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from parallel_process import parallel_process\n",
    "frames = parallel_process(files, get_edge_from_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_frame = pd.concat(frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_frame.to_csv('/rapids/notebooks/srabhi/champs-2019/input/parquet/baseline_edge_frame.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
